import torch
import os
import numpy as np
from PIL import Image


def singleton(cls):
    _instance = {}

    def inner():
        if cls not in _instance:
            _instance[cls] = cls()
        return _instance[cls]
    return inner

@singleton
class MidVisualizer():
    label = None
    def __init__(self):
        super().__init__()
        self.subject = ''
        self.on = False

    def set_save_dir(self,path):
        """
        存的大目录
        """
        if not self.on:
            return
        self.save_dir = path
        if not os.path.exists(path):
            os.makedirs(path)

    def set_save_subject(self,subject):
        """
        子文件夹名字，存的图代表的意义
        """
        if not self.on:
            return
        self.subject = subject
        curr_dir = os.path.join(self.save_dir,self.subject)
        if not os.path.exists(curr_dir):
            os.makedirs(curr_dir)
        self.save_path = os.path.join(self.save_dir,self.subject,self.name)

    def set_name(self,name):
        """
        存的文件名
        """
        if not self.on:
            return
        self.name = name
        assert self.save_dir is not None
        self.save_path = os.path.join(self.save_dir,self.subject,name)

    def set_target(self,t):
        """
`       t:tensor[b,c,h,w]
        
        """
        if not self.on:
            return
        self.squeezed_t = torch.squeeze(torch.sum(t, dim=1))

    def save_gray(self):
        """
        存成灰度图，按通道叠加，归一化
        """
        if not self.on:
            return
        gray = self.squeezed_t.detach().cpu().numpy()

        # normalize
        gray = (gray-np.min(gray))/(np.max(gray)-np.min(gray))*255

        Image.fromarray((gray).astype(np.uint8)).save(self.save_path)

    def set_on(self):
        self.on = True

if __name__ =='__main__':
    c1 = MidVisualizer()
    c2 = MidVisualizer()
    print(c1 is c2)